import tensorflow as tf
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from nets import nets_factory


def read_tfrecord(file_path):
    queue = tf.train.string_input_producer([file_path])
    _, serialized_example = tf.TFRecordReader().read(queue)
    features = tf.parse_single_example(serialized_example, features={
        "image": tf.FixedLenFeature([], dtype=tf.string),
        'label0': tf.FixedLenFeature([], dtype=tf.int64),
        'label1': tf.FixedLenFeature([], dtype=tf.int64),
        'label2': tf.FixedLenFeature([], dtype=tf.int64),
        'label3': tf.FixedLenFeature([], dtype=tf.int64),
    })

    # 获取图片数据
    image = tf.decode_raw(features['image'], tf.uint8)
    # tf.train.shuffle_batch必须确定shape

    # 没有经过预处理的灰度图
    image_raw = tf.reshape(image, [224, 224])
    # tf.train.shuffle_batch必须确定shape
    image = tf.reshape(image, [224, 224])

    # 图片预处理  图片的处理，将图片转成[-1,1]区间
    image = tf.cast(image, tf.float32) / 255.0
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)

    label0 = tf.cast(features['label0'], tf.int64)
    label1 = tf.cast(features['label1'], tf.int64)
    label2 = tf.cast(features['label2'], tf.int64)
    label3 = tf.cast(features['label3'], tf.int64)

    return image, image_raw, label0, label1, label2, label3


TFRECORD_FILE = 'F:/Idea workspace/tensorflow-learning/data/sample/tfrecord/captcha-test.tfrecord'

image, image_raw, label0, label1, label2, label3 = read_tfrecord(TFRECORD_FILE)

# 使用shuffle_batch可以随机打乱
image_batch, image_raw_batch, label_batch0, label_batch1, label_batch2, label_batch3 = tf.train.shuffle_batch(
    [image, image_raw, label0, label1, label2, label3], batch_size=1,
    capacity=50000, min_after_dequeue=10000, num_threads=1)

net_construct = nets_factory.get_network_fn('alexnet_v2', 10, weight_decay=0.005)

BATCH_SIZE = 1

MODEL_PATH = 'F:/Idea workspace/tensorflow-learning/data/model/crack_captcha.model-6000'

# placeholder
x = tf.placeholder(tf.float32, [None, 224, 224])

with tf.Session() as session:
    # inputs: a tensor of size [batch_size, height, width, channels]
    X = tf.reshape(x, [BATCH_SIZE, 224, 224, 1])
    net0, net1, net2, net3, end_points = net_construct(X)

    predict0 = tf.argmax(tf.reshape(net0, [1, 10]), 1)
    predict1 = tf.argmax(tf.reshape(net1, [1, 10]), 1)
    predict2 = tf.argmax(tf.reshape(net2, [1, 10]), 1)
    predict3 = tf.argmax(tf.reshape(net3, [1, 10]), 1)

    session.run(tf.global_variables_initializer())

    saver = tf.train.Saver()
    saver.restore(session, MODEL_PATH)

    # 创建一个协调器，管理线程
    coord = tf.train.Coordinator()

    # 启动QueueRunner, 此时文件名队列已经进队
    threads = tf.train.start_queue_runners(sess=session, coord=coord)

    for i in range(100):
        # 获取一个批次的数据和标签
        b_image, b_image_raw, b_label0, b_label1, b_label2, b_label3 = session.run([image_batch,
                                                                                    image_raw_batch,
                                                                                    label_batch0,
                                                                                    label_batch1,
                                                                                    label_batch2,
                                                                                    label_batch3])
        # 打印标签
        print('label:', b_label0, b_label1, b_label2, b_label3)
        # 预测
        label0, label1, label2, label3 = session.run([predict0, predict1, predict2, predict3], feed_dict={x: b_image})
        # 打印预测值
        print('predict:', label0, label1, label2, label3)
        print('-------------------------------------------------------')

        # 通知其他线程关闭
    coord.request_stop()
    # 其他所有线程关闭之后，这一函数才能返回
    coord.join(threads)
